package com.sjc.wordcount.batch;

import org.apache.flink.api.common.JobExecutionResult;
import org.apache.flink.api.common.accumulators.IntCounter;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.util.Collector;

/**
 * 统计单词词频
 * @author: sangjicun
 * @date: 2021/3/16 20:13
 */
public class WordCount {

    public static void main(String[] args) throws Exception {
        // 获取命令行参数
        final ParameterTool params = ParameterTool.fromArgs(args);
        // 获取一个执行环境，本地或者集群环境会自动识别
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.getConfig().setGlobalJobParameters(params);

        // 读取输入数据
        DataSet<String> text;
        if (params.has("input")) {
            // 读取text文件
            text = env.readTextFile(params.get("input"));
        } else {
            // 读取默认测试数据集
            System.out.println("Executing WordCount example with default input data set.");
            System.out.println("Use --input to specify file input.");
            text = WordCountData.getDefaultTextLineDataSet(env);
        }


        DataSet<Tuple2<String, Integer>> counts =
                text.flatMap(new Tokenizer())   // 切分每行单词
                        .groupBy(0)    // 对每个单词分组统计词频
                        .sum(1);

        if (params.has("output")) {
            // 数据输出为CSV格式
            counts.writeAsCsv(params.get("output"), "\n", " ");
            // 提交执行flink应用
            JobExecutionResult jobExecutionResult =env.execute("WordCount Example");
            jobExecutionResult.getAccumulatorResult("num-lines").toString();
        } else {
            System.out.println("Printing result to stdout. Use --output to specify output path.");
            // 数据打印控制台，内部封装了execute提交的flink应用
            counts.print();
        }

    }

    public static final class Tokenizer extends RichFlatMapFunction<String, Tuple2<String, Integer>> {
        private IntCounter numLines = new IntCounter();

        @Override
        public void open(Configuration parameters) throws Exception {
            getRuntimeContext().addAccumulator("num-lines", this.numLines);
            super.open(parameters);
        }

        @Override
        public void close() throws Exception {
            super.close();
        }

        @Override
        public void flatMap(String value, Collector<Tuple2<String, Integer>> out) {
            this.numLines.add(1);
            String[] tokens = value.toLowerCase().split("\\W+");

            for (String token : tokens) {
                if (token.length() > 0) {
                    out.collect(new Tuple2<>(token, 1));
                }
            }
        }
    }

}